from .base import PoseEstimator, center_ellipse, ransac_pnp, pose_calculate
import numpy as np
import cv2


class Cone(PoseEstimator):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def __name__(self):
        return "cone"

    def forward(self, C_3d, C_2d,*args, **ignore_kwargs):
        x_cnt = []
        X_cnt = []
        for c_3d, c_2d in zip(C_3d, C_2d):
            # 求解相机椭圆中心：
            x_cnt.append(center_ellipse(c_2d))
            X_cnt.append(center_ellipse(c_3d))
        # 求解相机位姿
        X_cnt = np.pad(X_cnt, ((0, 0), (0, 1)), "constant", constant_values=0)
        x_cnt = np.array(x_cnt)
        ind = ransac_pnp(self.K, X_cnt, x_cnt)
        if ind is None:
            return False, None, None
        R, T = pose_calculate(self.K, X_cnt[ind], x_cnt[ind])
        # nc = R @ np.array([0, 0, 1])
        # 批量求解新的陨石坑中心
        C_2d = np.array(C_2d)
        cnt = np.linalg.inv(C_2d @ self.K).transpose(0, 2, 1) @ R[:, 2]
        x_cnt_new = cnt[:, :2] / cnt[:, 2, None]
        ind = ransac_pnp(self.K, X_cnt, x_cnt_new)
        R, T = pose_calculate(self.K, X_cnt[ind], x_cnt_new[ind])
        if R is None:
            return False, None, None
        return True, R, T

    #     # 子范围内采样多个点，随机组成对应关系
    #     x_repeat = []
    #     X_repeat = []
    #     for i in range(10):
    #         x_repeat.append(
    #             x_cnt_new[ind] + np.random.randn(*x_cnt_new[ind].shape).clip(-1, 1)
    #         )
    #         X_repeat.append(X_cnt[ind])
    #     x_repeat = np.concatenate(x_repeat, axis=0)
    #     X_repeat = np.concatenate(X_repeat, axis=0)
    #     ind = ransac_pnp(self.K, X_repeat, x_repeat)
    #     if ind is None:
    #         return False, None, None
    #     R, T = pose_calculate(self.K, X_repeat[ind], x_repeat[ind])
    #     if R is None:
    #         return False, None, None
    #     return True, R, T

    # def vertical_ransac_pnp(K, points_3d, observations, th=8):
    #     best_num = 0
    #     best_inliers = np.array([False] * len(observations))
    #     sample = np.random.choice(len(observations), 4, replace=False)
    #     for _ in range(100):
    #         R, tvec = pose_calculate(K, points_3d[sample], observations[sample])
    #         if R is not None:
    #             err = np.linalg.norm(
    #                 cv2.projectPoints(points_3d, R, tvec, K, np.zeros(4))[0].squeeze()
    #                 - observations,
    #                 axis=1,
    #             )
    #             inliers = err < th
    #             if inliers.sum() > best_num:
    #                 best_num = inliers.sum()
    #                 best_inliers = inliers
    #             if inliers.sum() > 0.8 * len(observations):
    #                 break
    #             # 另一种方式，使用绝对点个数，实际上超过20个点后计算所得精度已经相当高了。
    #             # if inliers.sum() > 20:
    #             #     break
    #             if inliers.sum() > 4:
    #                 sample = np.random.choice(np.where(inliers)[0], 4, replace=False)
    #             else:
    #                 sample = np.random.choice(len(observations), 4, replace=False)
    #         else:
    #             sample = np.random.choice(len(observations), 4, replace=False)
    #     if best_inliers.sum() < 4:
    #         return None
    #     else:
    #         return best_inliers
